"""
IMAP连接池 - 核心优化模块
功能：复用IMAP连接，避免频繁建立TCP连接
"""

import imaplib
import threading
import time
from datetime import datetime
from contextlib import contextmanager
from queue import Queue, Empty
import logging

logger = logging.getLogger('email_system')


class IMAPConnection:
    """IMAP连接包装类"""
    
    def __init__(self, mail, account_id, created_at):
        self.mail = mail
        self.account_id = account_id
        self.created_at = created_at
        self.last_used = time.time()
        self.use_count = 0
    
    def is_alive(self):
        """检查连接是否存活 - 严格版（检查IMAP状态）"""
        try:
            # 1. 检查socket是否存在
            if not hasattr(self.mail, 'sock') or self.mail.sock is None:
                logger.debug(f"[连接池] 连接socket不存在")
                return False
            
            # 2. 检查socket状态
            import socket
            if self.mail.sock.fileno() == -1:
                logger.debug(f"[连接池] 连接socket已关闭")
                return False
            
            # 3. 检查IMAP层的state（核心！！）
            # imaplib有state属性，标准值包括 'NONAUTH', 'AUTH', 'SELECTED', 'LOGOUT'
            imap_state = getattr(self.mail, 'state', None)
            if imap_state not in ('AUTH', 'SELECTED'):
                logger.debug(f"[连接池] IMAP状态无效: {imap_state}")
                return False
            
            # 4. noop命令，确认服务端存活
            original_timeout = self.mail.sock.gettimeout()
            self.mail.sock.settimeout(3)  # 3秒超时
            
            try:
                status, _ = self.mail.noop()
                if status != 'OK':
                    logger.debug(f"[连接池] noop命令返回状态: {status}")
                    return False
            finally:
                # 恢复原始超时设置
                self.mail.sock.settimeout(original_timeout)
            
            logger.debug(f"[连接池] 连接状态检查通过 - IMAP状态: {imap_state}")
            return True
                
        except Exception as e:
            logger.debug(f"[连接池] 连接存活检查失败: {e}")
            return False
    
    def close(self):
        """关闭连接"""
        try:
            self.mail.close()
        except:
            pass
        try:
            self.mail.logout()
        except:
            pass


class IMAPConnectionPool:
    """IMAP连接池管理器"""
    
    def __init__(self, max_connections_per_account=3, idle_timeout=120, max_lifetime=600):
        """
        初始化连接池
        
        Args:
            max_connections_per_account: 每个账号最大连接数
            idle_timeout: 空闲超时时间（秒）
            max_lifetime: 连接最大生存时间（秒）
        """
        self.max_connections_per_account = max_connections_per_account
        self.idle_timeout = idle_timeout
        self.max_lifetime = max_lifetime
        
        # 连接池：{account_id: Queue of IMAPConnection}
        self.pools = {}
        self.locks = {}  # 每个账号的锁
        
        # 统计信息
        self.stats = {
            'total_created': 0,
            'total_reused': 0,
            'total_closed': 0,
        }
        
        # 启动清理线程
        self.running = True
        self.cleanup_thread = threading.Thread(target=self._cleanup_loop, daemon=True)
        self.cleanup_thread.start()
        
        logger.info("[连接池] 初始化完成")
    
    def _get_pool_and_lock(self, account_id):
        """获取或创建账号的连接池和锁"""
        if account_id not in self.pools:
            self.pools[account_id] = Queue(maxsize=self.max_connections_per_account)
            self.locks[account_id] = threading.Lock()
        
        return self.pools[account_id], self.locks[account_id]
    
    def _create_connection(self, account):
        """创建新的IMAP连接"""
        from utils.encryption import email_encryption
        
        logger.info(f"[连接池] 创建新连接 - 账号: {account.email}")
        start_time = time.time()
        
        try:
            password = email_encryption.decrypt_password(account.password)
            
            # 创建IMAP连接
            if account.use_ssl:
                mail = imaplib.IMAP4_SSL(
                    account.imap_server,
                    account.imap_port,
                    timeout=30
                )
            else:
                mail = imaplib.IMAP4(
                    account.imap_server,
                    account.imap_port,
                    timeout=30
                )
            
            # 登录
            mail.login(account.username, password)
            
            duration = time.time() - start_time
            self.stats['total_created'] += 1
            
            logger.info(f"[连接池] 连接创建成功 - 耗时: {duration:.2f}秒, 总创建数: {self.stats['total_created']}")
            
            return IMAPConnection(mail, account.id, time.time())
        
        except Exception as e:
            logger.error(f"[连接池] 创建连接失败: {e}")
            raise
    
    @contextmanager
    def get_connection(self, account):
        """
        获取连接的上下文管理器
        
        用法:
            with connection_pool.get_connection(account) as mail:
                mail.select('INBOX')
                # ... 使用mail
        """
        pool, lock = self._get_pool_and_lock(account.id)
        conn = None
        reused = False
        
        try:
            # 尝试从池中获取连接
            try:
                conn = pool.get_nowait()
                
                # 检查连接是否有效 - 优化版
                current_time = time.time()
                time_since_creation = current_time - conn.created_at
                time_since_last_use = current_time - conn.last_used
                
                # QQ邮箱特殊处理：完全禁用连接复用，每次都创建新连接
                logger.info(f"[连接池] QQ邮箱禁用连接复用，关闭现有连接 - 账号: {account.email}")
                conn.close()
                conn = self._create_connection(account)
            
            except Empty:
                # 池中没有连接，创建新的
                conn = self._create_connection(account)
            
            # 更新最后使用时间
            conn.last_used = time.time()
            
            # 返回原始mail对象供使用
            yield conn.mail
        
        finally:
            # QQ邮箱特殊处理：不归还连接，每次都关闭
            if conn:
                logger.info(f"[连接池] QQ邮箱不归还连接，直接关闭 - 账号: {account.email}")
                conn.close()
                self.stats['total_closed'] += 1
    
    def _cleanup_loop(self):
        """清理空闲连接的后台线程"""
        while self.running:
            try:
                time.sleep(30)  # 每30秒检查一次
                self._cleanup_idle_connections()
            except Exception as e:
                logger.error(f"[连接池] 清理线程异常: {e}")
    
    def _cleanup_idle_connections(self):
        """清理空闲连接"""
        current_time = time.time()
        total_cleaned = 0
        
        for account_id, pool in list(self.pools.items()):
            connections_to_keep = []
            
            # 检查所有连接
            while True:
                try:
                    conn = pool.get_nowait()
                    
                    # 检查是否超过空闲时间或最大生存时间
                    idle_time = current_time - conn.last_used
                    lifetime = current_time - conn.created_at
                    
                    if idle_time > self.idle_timeout or lifetime > self.max_lifetime:
                        logger.info(
                            f"[连接池] 清理连接 - 账号ID: {account_id}, "
                            f"空闲时间: {idle_time:.0f}秒, 生存时间: {lifetime:.0f}秒"
                        )
                        conn.close()
                        total_cleaned += 1
                        self.stats['total_closed'] += 1
                    else:
                        # 保留连接
                        connections_to_keep.append(conn)
                
                except Empty:
                    break
            
            # 将保留的连接放回池中
            for conn in connections_to_keep:
                try:
                    pool.put_nowait(conn)
                except:
                    conn.close()
                    total_cleaned += 1
        
        if total_cleaned > 0:
            logger.info(f"[连接池] 清理完成 - 关闭了 {total_cleaned} 个连接")
    
    def close_account_connections(self, account_id):
        """关闭指定账号的所有连接"""
        if account_id not in self.pools:
            return
        
        pool = self.pools[account_id]
        closed_count = 0
        
        while True:
            try:
                conn = pool.get_nowait()
                conn.close()
                closed_count += 1
            except Empty:
                break
        
        logger.info(f"[连接池] 关闭账号 {account_id} 的 {closed_count} 个连接")
    
    def close_all(self):
        """关闭所有连接"""
        self.running = False
        
        for account_id in list(self.pools.keys()):
            self.close_account_connections(account_id)
        
        logger.info(f"[连接池] 已关闭所有连接")
    
    def get_stats(self):
        """获取统计信息"""
        active_connections = sum(pool.qsize() for pool in self.pools.values())
        
        return {
            **self.stats,
            'active_connections': active_connections,
            'accounts_count': len(self.pools),
        }


# 全局连接池实例 - 针对QQ邮箱的保守策略
imap_connection_pool = IMAPConnectionPool(
    max_connections_per_account=1,  # QQ邮箱只保持1个连接
    idle_timeout=10,  # 10秒空闲回收（QQ邮箱容易LOGOUT）
    max_lifetime=30   # 30秒最大生存时间（QQ邮箱连接不稳定）
)
